#include <iostream>
#include <vector>
#include <string>
#include <hip/hip_runtime.h>

#define HIP_CHECK(cmd) \
do { \
    hipError_t error = cmd; \
    if (error != hipSuccess) { \
        std::cerr << "HIP error: " << hipGetErrorString(error) \
                  << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
        exit(EXIT_FAILURE); \
    } \
} while(0)

__global__ void vectorAdd(const float* A, const float* B, float* C, int numElements) {
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < numElements) {
        C[i] = A[i] + B[i];
    }
}

void print_usage(const char* prog_name) {
    std::cerr << "Usage: " << prog_name << " [--free]\n"
              << "Options:\n"
              << "  --free  free memory for leak testing\n";
}

int main(int argc, char* argv[]) {
    // 解析命令行参数
    bool free_memory = false;
    if (argc > 1) {
        std::string arg(argv[1]);
        if (arg == "--free") {
            free_memory = true;
        } else {
            print_usage(argv[0]);
            return EXIT_FAILURE;
        }
    }

    // 初始化配置
    const size_t N = 1 << 20;  // 3M elements
    const size_t bytes = N * sizeof(float);

    //======================================
    // 内存分配阶段
    //======================================

    // 1. 常规设备内存
    float *d_A, *d_B, *d_C;
    HIP_CHECK(hipMalloc(&d_A, bytes));
    HIP_CHECK(hipMalloc(&d_B, bytes));
    HIP_CHECK(hipMalloc(&d_C, bytes));

    // 2. 锁页内存
    float *h_pinned;
    HIP_CHECK(hipHostMalloc(&h_pinned, bytes));

    // 3. 统一内存
    float *u_mem;
    HIP_CHECK(hipMallocManaged(&u_mem, bytes));

    // 4. 注册主机内存
    float *h_registered = new float[N];
    HIP_CHECK(hipHostRegister(h_registered, bytes, hipHostRegisterDefault));

    // 5. 内存池分配
    hipMemPool_t memPool;
    HIP_CHECK(hipDeviceGetMemPool(&memPool, 0));
    float *d_pool;
    HIP_CHECK(hipMallocFromPoolAsync(&d_pool, bytes, memPool, 0));

    // 6. 三维内存分配
    hipPitchedPtr pitchedPtr;
    hipExtent extent = make_hipExtent(256, 256, sizeof(float));
    HIP_CHECK(hipMalloc3D(&pitchedPtr, extent));

    //======================================
    // 数据操作阶段
    //======================================

    // 初始化主机数据
    std::vector<float> h_A(N, 1.0f), h_B(N, 2.0f), h_C(N);

    // 主机到设备拷贝
    HIP_CHECK(hipMemcpy(d_A, h_A.data(), bytes, hipMemcpyHostToDevice));
    HIP_CHECK(hipMemcpy(d_B, h_B.data(), bytes, hipMemcpyHostToDevice));

    // 执行内核
    const int blockSize = 256;
    const int gridSize = (N + blockSize - 1) / blockSize;
    vectorAdd<<<gridSize, blockSize>>>(d_A, d_B, d_C, N);
    HIP_CHECK(hipGetLastError());
    HIP_CHECK(hipDeviceSynchronize());

    // 设备到主机拷贝
    HIP_CHECK(hipMemcpy(h_C.data(), d_C, bytes, hipMemcpyDeviceToHost));

    // 验证结果
    for (size_t i = 0; i < N; ++i) {
        if (fabs(h_C[i] - 3.0f) > 1e-5) {
            std::cerr << "Verification failed at index " << i << std::endl;
            exit(EXIT_FAILURE);
        }
    }

    //======================================
    // 条件内存释放阶段
    //======================================
    if (free_memory) {
        // 释放常规设备内存
        HIP_CHECK(hipFree(d_A));
        HIP_CHECK(hipFree(d_B));
        HIP_CHECK(hipFree(d_C));

        // 释放锁页内存
        HIP_CHECK(hipHostFree(h_pinned));

        // 释放统一内存
        HIP_CHECK(hipFree(u_mem));

        // 释放注册内存
        HIP_CHECK(hipHostUnregister(h_registered));
        delete[] h_registered;

        // 释放内存池分配
        HIP_CHECK(hipFree(d_pool));

        // 释放三维内存
        HIP_CHECK(hipFree(pitchedPtr.ptr));

        std::cout << "All resources released successfully\n";
    } else {
        std::cout << "Memory intentionally kept allocated for testing\n";
    }

    return EXIT_SUCCESS;
}
// 上面用例可以再增加hipExtHostMalloc、hipMalloc3DArray、hipMallocArray、hipMallocPitch、hipMemAllocPitch、hipMallocMipmappedArray等接口的调用
